import torch
from torch.utils.data import DataLoader
from torch.optim import SGD

from dataLoad import load_data, MyDataset, printTestOut, getTestOutInfo
from model import CNNModel
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import os
#from model import ResNet18

os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'

lr = 0.005
num_epoches = 50
use_GPU = True
model_save_path = 'model_172919.pt'  # 模型保存路径


def loss_fn(logits, labels):
    loss_fn = F.binary_cross_entropy_with_logits(logits, labels)
    return loss_fn

def loss_fn_domain(logits, labels, weights, domain_use):
    loss_fn = F.binary_cross_entropy_with_logits(logits, labels)
    weighted_loss = loss_fn * weights * domain_use
    # print(domain_use)
    mean_weighted_loss = weighted_loss.mean()
    return mean_weighted_loss



## --------移动静止测试-------------
# 加载训练集
id_list = [1, 2, 3, 4, 5,6,7,8,9,10]
path1 = [('Gather_Data/LOS01_20250308/已筛选/LOS_CombinedData10.mat', 'CSI', id_list, False)]
path2= [('Gather_Data/NLOS01_20250309/已筛选/NLOS_CombinedData10.mat', 'CSI', id_list, True)]
# path1 = [('Gather_STA_CSI/LOS+NLOS/LOS1_20250105/AP01/LOS_CombinedData10.mat', 'CSI', id_list, False)]
# path2 = [('Gather_STA_CSI/LOS+NLOS/NLOS1_20250109/AP01/NLOS_CombinedData10.mat', 'CSI', id_list, True)]
train_data_pair, train_rff_label, train_domain_label, train_domain_use, dev_data_pair, dev_rff_label, dev_domain_label, dev_domain_use, id_list = load_data(
    data_list=[path2, path1], overlay_num=20)
train_dataset = MyDataset(train_data_pair, train_rff_label, train_domain_label, train_domain_use)

#加载测试集
id_list = [1, 2,3,4,5,6,7,8,9,10]
# path2 = [('Gather_STA_CSI/LOS+NLOS/NLOS2_20250109/AP01/NLOS_CombinedData10.mat', 'CSI', id_list, True)]
path1 = [('Gather_Data/LOS02_20250308/已筛选/LOS_CombinedData10.mat', 'CSI', id_list, False)]
# path2 = [('Gather_Data/NLOS02_20250309/已筛选/NLOS_CombinedData5.mat', 'CSI', id_list, True)]
# path2 = [('1029_rff_ddd_raw.mat', 'rff', id_list, False)]
# path1 = [('1101_rff_ddd_raw.mat', 'rff', id_list, True)]
test_data_pair, test_rff_label, test_domain_label, test_domain_use, dev_data_pair, dev_rff_label, dev_domain_label, dev_domain_use, id_list = load_data(
     data_list=[path1], overlay_num=20)

output_savepath = "out_fi_1029.txt"
dev_dataset = MyDataset(dev_data_pair, dev_rff_label, dev_domain_label, dev_domain_use)

DMRS_num = 1

hparams = {
    'DMRS_num': DMRS_num,
    'out_feature': 600
}

train_dataloader = DataLoader(train_dataset, batch_size=16, drop_last=True)

dev_dataloader = DataLoader(dev_dataset, batch_size=16, drop_last=True)

model = CNNModel(hparams)
# model = ResNet18(hparams)

sgd = SGD(model.parameters(), lr=float(lr))
num_epoches = int(num_epoches)

if use_GPU:
    device = torch.device("cuda")
else:
    device = torch.device("cpu")
model.to(device)

alpha = -0.5
max_acc = 0
epoch_train_loss_list = []
epoch_train_domain_loss_list = []
epoch_dev_loss_list = []
epoch_train_acc_list = []
epoch_dev_acc_list = []
epoch_dev_domain_loss_list = []
#对照实验：去除域LOSS，域标签分类
# for epoch in range(num_epoches):
#     epoch_train_loss = 0
#     epoch_train_acc = 0
#     epoch_dev_loss = 0
#     epoch_dev_acc = 0
#     train_num = 0
#     dev_num = 0
#     test_dict = {}
#
#     # 训练阶段
#     model.train()
#     for x_1, x_2, label_list, domain_label, domain_use in train_dataloader:
#         label = label_list[:, 0:1]
#         s = len(label)
#
#         sgd.zero_grad()
#         class_output, _, feature_x1, feature_x2 = model(x_1.to(device).float(), x_2.to(device).float(), alpha)  # 去掉 domain_output
#
#         # 只计算分类损失
#         loss = loss_fn(class_output, label.to(device))
#
#         epoch_train_loss += loss.item()
#         label_pred = torch.where(class_output > 0, torch.ones_like(class_output),
#                                  torch.zeros_like(class_output)).cpu()
#         acc = torch.sum(label_pred == label)
#         epoch_train_acc += acc.item() / s
#         train_num += 1
#         loss.backward()
#         sgd.step()
#
#     # 验证阶段
#     model.eval()
#     legal_total = 0  # 合法设备总数
#     legal_correct = 0  # 正确识别的合法设备数
#     illegal_total = 0  # 非法设备总数
#     illegal_correct = 0  # 正确拒绝的非法设备数
#     with torch.no_grad():
#         test_dict = {}
#         for x_1, x_2, label_list, domain_label, domain_use in dev_dataloader:
#             label = label_list[:, 0:1]
#             s = len(label)
#
#             class_output, _, feature_x1, feature_x2 = model(x_1.to(device).float(), x_2.to(device).float(), alpha)  # 去掉 domain_output
#
#             # 只计算分类损失
#             loss = loss_fn(class_output, label.to(device))
#
#             epoch_dev_loss += loss.item()
#             label_pred = torch.where(class_output > 0, torch.ones_like(class_output),
#                                      torch.zeros_like(class_output)).cpu()
#
#             test_dict = getTestOutInfo(label_pred, label_list, id_list, test_dict)
#             acc = torch.sum(label_pred == label)
#             epoch_dev_acc += acc.item() / s
#             dev_num += 1
#
#             # 统计合法设备和非法设备的识别情况
#             legal_indices = (label == 1).nonzero(as_tuple=True)[0]
#             illegal_indices = (label == 0).nonzero(as_tuple=True)[0]
#
#             legal_total += len(legal_indices)
#             legal_correct += torch.sum(label_pred[legal_indices] == label[legal_indices]).item()
#
#             illegal_total += len(illegal_indices)
#             illegal_correct += torch.sum(label_pred[illegal_indices] == label[illegal_indices]).item()
#
#         printTestOut(test_dict, save_path=output_savepath)
#
#     # 计算合法设备准确率和非法设备拒绝率
#     if legal_total > 0:
#         legal_accuracy = legal_correct / legal_total
#     else:
#         legal_accuracy = 0
#
#     if illegal_total > 0:
#         illegal_rejection_rate = illegal_correct / illegal_total
#     else:
#         illegal_rejection_rate = 0
#
#     # 记录训练和验证的损失、准确率等指标
#     if train_num > 0:
#         epoch_train_loss_list.append(epoch_train_loss / train_num)
#         epoch_train_acc_list.append(epoch_train_acc / train_num)
#     if dev_num > 0:
#         epoch_dev_loss_list.append(epoch_dev_loss / dev_num)
#         epoch_dev_acc_list.append(epoch_dev_acc / dev_num)
#
#     dev_acc = epoch_dev_acc / dev_num if dev_num > 0 else 0
#     dev_loss = epoch_dev_loss / dev_num if dev_num > 0 else 0
#     if dev_acc > max_acc:
#         print('save')
#         max_acc = dev_acc
#         torch.save(model, model_save_path)
#     print("epoch {:.4f} train acc: {:.4f},train loss: {:.4f}, dev acc: {:.4f}, dev loss: {:.4f}".format(epoch,
#                                                                                                         epoch_train_acc / train_num if train_num > 0 else 0,
#                                                                                                         epoch_train_loss / train_num if train_num > 0 else 0,
#                                                                                                         epoch_dev_acc / dev_num if dev_num > 0 else 0,
#                                                                                                         epoch_dev_loss / dev_num if dev_num > 0 else 0))
#     print(f"合法设备准确率: {legal_accuracy:.4f}, 非法设备拒绝率: {illegal_rejection_rate:.4f}")


for epoch in range(num_epoches):
    epoch_train_loss = 0
    epoch_train_domain_loss = 0
    epoch_train_acc = 0
    epoch_dev_loss = 0
    epoch_dev_acc = 0
    epoch_dev_domain_loss = 0
    train_num = 0
    dev_num = 0
    test_dict = {}
    for x_1, x_2, label_list, domain_label, domain_use in train_dataloader:
        label = label_list[:, 0:1]
        label_domain = domain_label[:, 0:1]
        s = len(label)
        weights = torch.randn(s)

        sgd.zero_grad()
        class_output, domain_output, feature_x1, feature_x2 = model(x_1.to(device).float(), x_2.to(device).float(),
                                                                    alpha)

        # 计算原损失
        loss = loss_fn(class_output, label.to(device)) + loss_fn_domain(domain_output, label_domain.to(device),
                                                                        label.to(device), domain_use.to(device))

        epoch_train_loss += loss.item()
        epoch_train_domain_loss += loss_fn_domain(domain_output, label_domain.to(device), label.to(device),
                                                  domain_use.to(device)).item()
        label_pred = torch.where(class_output > 0, torch.ones_like(class_output),
                                 torch.zeros_like(class_output)).cpu()
        acc = torch.sum(label_pred == label)
        epoch_train_acc += acc.item() / s
        train_num += 1
        loss.backward()
        sgd.step()

    model.eval()
    legal_total = 0  # 合法设备总数
    legal_correct = 0  # 正确识别的合法设备数
    illegal_total = 0  # 非法设备总数
    illegal_correct = 0  # 正确拒绝的非法设备数
    with torch.no_grad():
        test_dict = {}
        for x_1, x_2, label_list, domain_label, domain_use in dev_dataloader:
            label = label_list[:, 0:1]
            label_domain = domain_label[:, 0:1]
            s = len(label)
            weights = torch.randn(s)

            sgd.zero_grad()
            class_output, domain_output, feature_x1, feature_x2 = model(x_1.to(device).float(), x_2.to(device).float(),
                                                                        alpha)
            # 计算原损失
            loss = loss_fn(class_output, label.to(device)) + loss_fn_domain(domain_output, label_domain.to(device),
                                                                            label.to(device), domain_use.to(device))

            epoch_dev_loss += loss.item()
            epoch_dev_domain_loss += loss_fn_domain(domain_output, label_domain.to(device), label.to(device),
                                                    domain_use.to(device)).item()
            label_pred = torch.where(class_output > 0, torch.ones_like(class_output),
                                     torch.zeros_like(class_output)).cpu()

            test_dict = getTestOutInfo(label_pred, label_list, id_list, test_dict)
            acc = torch.sum(label_pred == label)
            epoch_dev_acc += acc.item() / s
            dev_num += 1

            # 假设 1 为合法设备，0 为非法设备
            legal_indices = (label == 1).nonzero(as_tuple=True)[0]
            illegal_indices = (label == 0).nonzero(as_tuple=True)[0]

            legal_total += len(legal_indices)
            legal_correct += torch.sum(label_pred[legal_indices] == label[legal_indices]).item()

            illegal_total += len(illegal_indices)
            illegal_correct += torch.sum(label_pred[illegal_indices] == label[illegal_indices]).item()

        printTestOut(test_dict, save_path=output_savepath)

    # 计算合法设备准确率和非法设备拒绝率
    if legal_total > 0:
        legal_accuracy = legal_correct / legal_total
    else:
        legal_accuracy = 0

    if illegal_total > 0:
        illegal_rejection_rate = illegal_correct / illegal_total
    else:
        illegal_rejection_rate = 0

    epoch_train_loss_list.append(epoch_train_loss / train_num)
    epoch_train_domain_loss_list.append(epoch_train_domain_loss / train_num)
    epoch_dev_domain_loss_list.append(epoch_dev_domain_loss / train_num)
    epoch_dev_loss_list.append(epoch_dev_loss / dev_num)
    epoch_train_acc_list.append(epoch_train_acc / train_num)
    epoch_dev_acc_list.append(epoch_dev_acc / dev_num)
    dev_acc = epoch_dev_acc / dev_num
    dev_loss = epoch_dev_loss / dev_num
    if dev_acc > max_acc:
        print('save')
        max_acc = dev_acc
        torch.save(model, model_save_path)
    print("epoch {:.4f} train acc: {:.4f},train loss: {:.4f}, dev acc: {:.4f}, dev loss: {:.4f}".format(epoch,
                                                                                                        epoch_train_acc / train_num,
                                                                                                        epoch_train_loss / train_num,
                                                                                                        epoch_dev_acc / dev_num,
                                                                                                        epoch_dev_loss / dev_num))
    print(f"合法设备准确率: {legal_accuracy:.4f}, 非法设备拒绝率: {illegal_rejection_rate:.4f}")

#
# #重新设置了预测逻辑
# for epoch in range(num_epoches):
#     epoch_train_loss = 0
#     epoch_train_domain_loss = 0
#     epoch_train_acc = 0
#     epoch_dev_loss = 0
#     epoch_dev_acc = 0
#     epoch_dev_domain_loss = 0
#     train_num = 0
#     dev_num = 0
#     test_dict = {}
#
#     # 训练阶段
#     model.train()
#     for x_1, x_2, label_list, domain_label, domain_use in train_dataloader:
#         label = label_list[:, 0:1]
#         label_domain = domain_label[:, 0:1]
#         s = len(label)
#         weights = torch.randn(s)
#
#         sgd.zero_grad()
#         class_output, domain_output, feature_x1, feature_x2 = model(x_1.to(device).float(), x_2.to(device).float(),
#                                                                    alpha)
#
#         # 打印训练阶段的 class_output
#         print(f"训练阶段 - 当前批次的 class_output: {class_output.cpu().detach().numpy()}")
#
#         # 计算原损失
#         loss = loss_fn(class_output, label.to(device)) + loss_fn_domain(domain_output, label_domain.to(device),
#                                                                         label.to(device), domain_use.to(device))
#
#         epoch_train_loss += loss.item()
#         epoch_train_domain_loss += loss_fn_domain(domain_output, label_domain.to(device), label.to(device),
#                                                   domain_use.to(device)).item()
#
#         # 根据 class_output 进行阈值判断
#         legal_pred = (class_output >= 1) & (class_output <= 10)
#         illegal_pred = (class_output >= 11) & (class_output <= 15)
#         label_pred = torch.where(legal_pred, torch.ones_like(class_output),
#                                  torch.where(illegal_pred, torch.zeros_like(class_output),
#                                              torch.zeros_like(class_output)))
#         label_pred = label_pred.cpu()
#
#         acc = torch.sum(label_pred == label)
#         epoch_train_acc += acc.item() / s
#         train_num += 1
#         loss.backward()
#         sgd.step()
#
#     # 验证阶段
#     model.eval()
#     legal_total = 0  # 合法设备总数
#     legal_correct = 0  # 正确识别的合法设备数
#     illegal_total = 0  # 非法设备总数
#     illegal_correct = 0  # 正确拒绝的非法设备数
#     with torch.no_grad():
#         test_dict = {}
#         for x_1, x_2, label_list, domain_label, domain_use in dev_dataloader:
#             label = label_list[:, 0:1]
#             label_domain = domain_label[:, 0:1]
#             s = len(label)
#             weights = torch.randn(s)
#
#             sgd.zero_grad()
#             class_output, domain_output, feature_x1, feature_x2 = model(x_1.to(device).float(), x_2.to(device).float(),
#                                                                        alpha)
#
#             # 打印验证阶段的 class_output
#             print(f"验证阶段 - 当前批次的 class_output: {class_output.cpu().detach().numpy()}")
#
#             # 计算原损失
#             loss = loss_fn(class_output, label.to(device)) + loss_fn_domain(domain_output, label_domain.to(device),
#                                                                             label.to(device), domain_use.to(device))
#
#             epoch_dev_loss += loss.item()
#             epoch_dev_domain_loss += loss_fn_domain(domain_output, label_domain.to(device), label.to(device),
#                                                     domain_use.to(device)).item()
#
#             # 根据 class_output 进行阈值判断
#             legal_pred = (class_output >= 1) & (class_output <= 10)
#             illegal_pred = (class_output >= 11) & (class_output <= 15)
#             label_pred = torch.where(legal_pred, torch.ones_like(class_output),
#                                      torch.where(illegal_pred, torch.zeros_like(class_output),
#                                                  torch.zeros_like(class_output)))
#             label_pred = label_pred.cpu()
#
#             # 统计合法和非法设备的识别情况
#             for i in range(len(label_pred)):
#                 if int(label[i].item()) == 1:
#                     legal_total += 1
#                     if label_pred[i] == 1:
#                         legal_correct += 1
#                 elif int(label[i].item()) == 0:
#                     illegal_total += 1
#                     if label_pred[i] == 0:
#                         illegal_correct += 1
#
#             test_dict = getTestOutInfo(label_pred, label_list, id_list, test_dict)
#             acc = torch.sum(label_pred == label)
#             epoch_dev_acc += acc.item() / s
#             dev_num += 1
#
#         printTestOut(test_dict, save_path=output_savepath)
#
#     # 计算合法设备准确率和非法设备拒绝率
#     if legal_total > 0:
#         legal_accuracy = legal_correct / legal_total
#     else:
#         legal_accuracy = 0
#
#     if illegal_total > 0:
#         illegal_rejection_rate = illegal_correct / illegal_total
#     else:
#         illegal_rejection_rate = 0
#
#     epoch_train_loss_list.append(epoch_train_loss / train_num)
#     epoch_train_domain_loss_list.append(epoch_train_domain_loss / train_num)
#     epoch_dev_domain_loss_list.append(epoch_dev_domain_loss / train_num)
#     epoch_dev_loss_list.append(epoch_dev_loss / dev_num)
#     epoch_train_acc_list.append(epoch_train_acc / train_num)
#     epoch_dev_acc_list.append(epoch_dev_acc / dev_num)
#     dev_acc = epoch_dev_acc / dev_num
#     dev_loss = epoch_dev_loss / dev_num
#     if dev_acc > max_acc:
#         print('save')
#         max_acc = dev_acc
#         torch.save(model, model_save_path)
#     print("epoch {:.4f} train acc: {:.4f},train loss: {:.4f}, dev acc: {:.4f}, dev loss: {:.4f}".format(epoch,
#                                                                                                         epoch_train_acc / train_num,
#                                                                                                         epoch_train_loss / train_num,
#                                                                                                         epoch_dev_acc / dev_num,
#                                                                                                         epoch_dev_loss / dev_num))
#     print(f"合法设备准确率: {legal_accuracy:.4f}, 非法设备拒绝率: {illegal_rejection_rate:.4f}")
